/*=============================================================
* 	Copyright (C) 2024 USTC All rights reserved.
* 	
* 	filename : core.h
* 	Author: Wei Jinsong
* 	Mail: vitus@mainl.ustc.edu.cn
* 	Date 2024-05-13
* 	Abstract 
* 		
*
=============================================================*/

#ifndef __CORE_H__
#define __CORE_H__

#include <iostream>
#include <string>
#include <vector>

#include "systemc.h"
#include "common.h"
#include "tlm.h"
#include "tlm_utils/simple_initiator_socket.h"
#include "tlm_utils/simple_target_socket.h"
#include "nlohmann/json.hpp"

using json = nlohmann::json;
using namespace sc_core;

class BaseCore 
    : public::sc_module
    , public::SpikeInterface<>
{
public :
    SC_HAS_PROCESS(BaseCore);
    BaseCore(sc_module_name name, uint32_t mID, sc_clock& clock, string file);
            // , uint32_t recv_pipeline, uint32_t send_pipeline);
    ~BaseCore();

    sc_in_clk clk;
    SpikePayload inSpikeSig;
    SpikePayload outSpikeSig;

    sc_time clkPeriod;
    sc_port<SpikeInterface<>> txSpikePort;
    sc_port<SynapseInterface<SpikePayload>> synapse;

public :
    const unsigned int mID;
    string type = "output";
    string label= "";
    map<uint32_t, uint32_t> _inSpikeCnt;
    map<uint32_t, uint32_t> _outSpikeCnt;

    vector<double> featureMaps;
    double featureTH;
    int featureType;
    int featureToken;

    sc_fifo<SpikePayload*> spikeFifo;
    SpikePayload* currSpike;
    sc_event recvEvent;

    uint32_t recvPipeline;
    uint32_t sendPipeline;
    SpikeEventQueue<SpikePayload> recvSpikePeq; 
    SpikeEventQueue<SpikePayload> sendSpikePeq;

    void sc_trace_spike(string name, SpikePayload& spike);

    virtual void generate_spike();
    virtual void recv_spike();
    virtual bool spike_transport_fw(SpikePayload& trans, sc_core::sc_time &delay);

    virtual void exec() = 0;
    virtual void set_params(string file) = 0;
    virtual void fs_exec() =0;
    // {
    //     // assert(0&&"Can't input the function");
    // };
    virtual void set_featureMap(string file);
    virtual void exec_input();

    virtual void print_spike_sta(ostream& os);
    int64_t get_send_spike();
    int64_t get_recv_spike();
};

class SpikeNeuronCore
    : public BaseCore
{
public :
    SC_HAS_PROCESS(SpikeNeuronCore);

    SpikeNeuronCore(string name, uint32_t mID, sc_clock& clock
            , string file, string infile); 
    ~SpikeNeuronCore();

private:
    double _t;
    double _th;
    // vector<double>_Vmem;
    uint32_t neuNum;

    // bool* _Spikes;
    // double* _Vmem;
    sc_signal<bool>* _Spikes;
    sc_signal<double>* _Vmem;

public :
    void set_params(string file) override;
    void exec() override;
    void fs_exec() override {}
};

class FSNeuronCore 
    : public BaseCore
{
public :
    SC_HAS_PROCESS(FSNeuronCore);
    FSNeuronCore(string name, uint32_t mId, sc_clock& clock
            , string file, string infile); 
    ~FSNeuronCore();
    struct ConvPoolParam{
        int ic=0;
        int oc=0;
        int kh=0;
        int kw=0;
        int sh=0;
        int sw=0;
        int ph=0;
        int pw=0;
    };
    struct FeatureMapParam{
        int channel;
        int height;
        int width;
    };
protected:
    uint32_t neuNum;
    ConvPoolParam _poolParam;
    ConvPoolParam _convParam;
    FeatureMapParam _featureParam;
    vector<double> biasParam;

    int fs_phase;
    int fs_k;
    vector<double> fs_h; // vreset
    vector<double> fs_d; // 
    vector<double> fs_t; // threshold
    vector<double> fs_tau;

    sc_signal<bool>     _fs_exec_state;
    sc_signal<bool>     _fs_recv_state;
    sc_signal<bool>     _switchReady;

    // uint32_t _currPhase;
    uint32_t _currStep;
    // uint32_t _currVmem;

    sc_signal<uint32_t> _currPhase;
    // sc_signal<uint32_t> _currStep;
    sc_signal<uint32_t> _currVmem;

    vector<double> _Vmem[2];
public:
    // sc_port<TimerInterface> timer;
    void set_params(string file) override;
    void set_featureMap(string file) override;
    void send_spike(SpikePayload&spike);
    void send_spike(int32_t srcneuid, int32_t phase, int32_t step, bool switch_phase, double value=1.0);
    virtual void print_spike_sta(ostream& os) override;

    virtual bool update_neuron(int32_t currVmem, int32_t currStep, int32_t neuid);
    virtual void exec_input() override;
    virtual void exec_vmem();
    virtual void exec() override;
    virtual void fs_exec() override;
    virtual void output_sta(SpikePayload *spike);
};

class AFSNeuronCore 
    : public FSNeuronCore
{
public:
    SC_HAS_PROCESS(AFSNeuronCore);
    AFSNeuronCore(string name, uint32_t mId, sc_clock& clock
            , string file, string infile); 
    ~AFSNeuronCore();

    map<int, int> _dvs_output;
    map<int, int> _dvs_output2;
    vector<double> _ExecVmem;
    virtual void exec_input() override;
    virtual bool update_neuron(int32_t currVmem, int32_t currStep, int32_t neuid) override;
    virtual void output_sta(SpikePayload *spike) override;
    virtual void print_spike_sta(ostream& os) override;
};

#endif

